///
/// Copyright (C) 2018, Cyberhaven
///
/// Permission is hereby granted, free of charge, to any person obtaining a copy
/// of this software and associated documentation files (the "Software"), to deal
/// in the Software without restriction, including without limitation the rights
/// to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
/// copies of the Software, and to permit persons to whom the Software is
/// furnished to do so, subject to the following conditions:
///
/// The above copyright notice and this permission notice shall be included in all
/// copies or substantial portions of the Software.
///
/// THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
/// IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
/// FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
/// AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
/// LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
/// OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
/// SOFTWARE.
///

#include <zlib.h>

#include <s2e/S2EExecutionState.h>

#include "PovGenerator.h"

namespace s2e {
namespace plugins {
namespace pov {

using namespace klee;

std::string PovGenerator::getFileName(S2EExecutionState *state, const PovOptions &opt, const std::string &filePrefix,
                                      const std::string &fileExtWithoutDot) {
    std::stringstream povFileNameSS;

    if (filePrefix.length() > 0) {
        povFileNameSS << filePrefix << "-";
    }

    switch (opt.m_type) {
        case POV_TYPE1:
            povFileNameSS << "pov-type1-";
            break;
        case POV_TYPE2:
            povFileNameSS << "pov-type2-";
            break;
        default:
            povFileNameSS << "pov-unknown-";
            break;
    }

    povFileNameSS << state->getID();

    if (fileExtWithoutDot.size() > 0) {
        povFileNameSS << "." << fileExtWithoutDot;
    }

    return g_s2e->getOutputFilename(povFileNameSS.str());
}

bool PovGenerator::writeToFile(const std::string &fileName, const void *data, unsigned size, bool compress) {
    std::vector<uint8_t> contents;
    if (compress) {
        contents = PovGenerator::compress(data, size);
        data = contents.data();
        size = contents.size();
    }

    FILE *povFile = fopen(fileName.c_str(), "wb");

    if (!povFile) {
        g_s2e->getWarningsStream() << "Could not create POV file" << '\n';
        return false;
    }

    if (fwrite(data, 1, size, povFile) != size) {
        int err = errno;
        g_s2e->getWarningsStream() << "Could not write POV file (" << strerror(err) << ")\n";
        return false;
    }

    if (fclose(povFile) != 0) {
        g_s2e->getWarningsStream() << "Error closing POV file" << '\n';
        return false;
    }

    g_s2e->getInfoStream() << "POV saved to file " << fileName << "\n";

    return true;
}

///
/// \brief PovGenerator::writeToFile save a PoV to a file.
/// PoV submission has to be carried out in a separate plugin.
///
/// \return the file name
///
std::string PovGenerator::writeToFile(S2EExecutionState *state, const PovOptions &opt, const std::string &filePrefix,
                                      const std::string &fileExtWithoutDot, const void *data, unsigned data_size) {

    std::string fileName = getFileName(state, opt, filePrefix, fileExtWithoutDot);
    if (m_compress) {
        fileName = fileName + ".zlib";
    }

    if (!writeToFile(fileName, data, data_size, m_compress)) {
        getWarningsStream(state) << "Could not write PoV " << fileName << "\n";
        exit(-1);
    }

    return fileName;
}

/* http://stackoverflow.com/questions/4538586/how-to-compress-a-buffer-with-zlib */
void PovGenerator::compress(const void *in_data, size_t in_data_size, std::vector<uint8_t> &out_data) {
    std::vector<uint8_t> buffer;

    const size_t BUFSIZE = 128 * 1024;
    uint8_t temp_buffer[BUFSIZE];

    z_stream strm;
    strm.zalloc = 0;
    strm.zfree = 0;
    strm.next_in = const_cast<uint8_t *>(reinterpret_cast<const uint8_t *>(in_data));
    strm.avail_in = in_data_size;
    strm.next_out = temp_buffer;
    strm.avail_out = BUFSIZE;

    deflateInit(&strm, Z_BEST_COMPRESSION);

    while (strm.avail_in != 0) {
        int res = deflate(&strm, Z_NO_FLUSH);
        s2e_assert(nullptr, res == Z_OK, "Deflate error");
        if (strm.avail_out == 0) {
            buffer.insert(buffer.end(), temp_buffer, temp_buffer + BUFSIZE);
            strm.next_out = temp_buffer;
            strm.avail_out = BUFSIZE;
        }
    }

    int deflate_res = Z_OK;
    while (deflate_res == Z_OK) {
        if (strm.avail_out == 0) {
            buffer.insert(buffer.end(), temp_buffer, temp_buffer + BUFSIZE);
            strm.next_out = temp_buffer;
            strm.avail_out = BUFSIZE;
        }
        deflate_res = deflate(&strm, Z_FINISH);
    }

    s2e_assert(nullptr, deflate_res == Z_STREAM_END, "Deflate error");
    buffer.insert(buffer.end(), temp_buffer, temp_buffer + BUFSIZE - strm.avail_out);
    deflateEnd(&strm);

    out_data.swap(buffer);
}

std::vector<uint8_t> PovGenerator::compress(const std::string &s) {
    return compress(s.c_str(), s.size());
}

std::vector<uint8_t> PovGenerator::compress(const void *data, unsigned size) {
    std::vector<uint8_t> compressed;
    compress(data, size, compressed);
    return compressed;
}

bool PovGenerator::solveConstraints(S2EExecutionState *state, const PovOptions &opt, Assignment &assignment) {
    using namespace klee;

    if (opt.m_extraConstraints.size() > 0) {
        ConstraintManager constraints = state->constraints();
        for (auto it : opt.m_extraConstraints) {
            constraints.addConstraint(it);
        }

        // TODO: all this constraints to assignment could be moved to klee
        // TODO: try to preserve existing concolics as much as possible

        // Extract symbolic objects
        ArrayVec symbObjects = state->symbolics;
        std::vector<std::vector<unsigned char>> concreteObjects;

        // XXX: Not sure about passing state, we may have new constraints
        Solver *solver = state->solver()->solver;
        Query query(constraints, ConstantExpr::alloc(0, Expr::Bool));
        if (!solver->getInitialValues(query, symbObjects, concreteObjects)) {
            getWarningsStream() << "Could not get symbolic solution\n";
            return false;
        }

        for (unsigned i = 0; i < symbObjects.size(); ++i) {
            assignment.add(symbObjects[i], concreteObjects[i]);
        }
    } else {
        assignment = *state->concolics;
    }

    return true;
}
} // namespace pov
} // namespace plugins
} // namespace s2e
